#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : util.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/14 13:10
import io
import asyncio
import time

from apps import logic
from apps.models import User
from apps import file_extract_util
from apps.utils import MinioUtil
from apps.utils.enumerate import KnowledgeFileStateEnum
from apps.vector_store_util import MilvusUtil
from core.setting import settings
from langchain.text_splitter import CharacterTextSplitter
from apps import embedding_utils


async def handle_file(file_id: int, knowledge_id: int, user: User) -> object:
    knowledge = await logic.knowledge_logic.knowledge_object(knowledge_id=knowledge_id, user=user)
    if not knowledge:
        print('没有找到知识库', f'knowledge_id: {knowledge_id}')
        return
    file = await logic.knowledge_file_logic.get_file_object(file_id=file_id, knowledge_id=knowledge_id)
    if not file:
        print('没有找到文件', f'file_id: {file_id}')
        return
    embedding_server = await logic.embedding_server_logic.embedding_server_info(
        embedding_server_id=knowledge.embedding_server_id)
    if not embedding_server:
        file.state = KnowledgeFileStateEnum.failed
        await file.save()
        return
        # 先将文件内容提取出来
    # 获取文件内容抽取服务
    extract = await logic.file_extract_logic.get_extract_object(extract_id=file.file_extract_id)
    if not extract:
        file.state = KnowledgeFileStateEnum.failed
        await file.save()
        return
    if extract.firm not in file_extract_util.CONFIG:
        file.state = KnowledgeFileStateEnum.failed
        await file.save()
        return
    file.state = KnowledgeFileStateEnum.processing
    print('文件内容提取加载', file_extract_util.CONFIG[extract.firm])
    extract_object = getattr(file_extract_util, file_extract_util.CONFIG[extract.firm])(api_key=extract.api_key)
    # 判断当前文件是否支持解析
    support_file_type = await extract_object.get_support_file_type()
    file_type = file.file_address.split('.')[-1]
    if file_type not in support_file_type:
        file.state = KnowledgeFileStateEnum.failed
        await file.save()
        return
    minio_util = MinioUtil(
        secret_key=settings.MINIO_SECRET_KEY,
        access_key=settings.MINIO_ACCESS_KEY,
        endpoint=settings.MINIO_ENDPOINT,
    )
    bucket_name = settings.MINIO_BASE_BUCKET
    object_name = file.file_address
    file_object = minio_util.get_object(bucket_name=bucket_name, object_name=object_name)
    # 将二进制数据包装为文件对象
    file_like_object = io.BytesIO(file_object)
    try:
        content = await extract_object.extract_file_content(base_url=extract.model_base_url,
                                                            file_object=(file.file_name, file_like_object))
    except:
        file.state = KnowledgeFileStateEnum.failed
        await file.save()
        return
    print('数据提取完成')
    # 进行数据切分处理
    text_splitter = CharacterTextSplitter(separator=file.split_str, chunk_size=file.split_size,
                                          chunk_overlap=int(file.split_size))
    print('数据切分完成')
    content = text_splitter.split_text(content)
    # 将数据转成向量
    if not embedding_utils.CONFIG.get(embedding_server.server_type):
        print('没有找到向量转换方法', embedding_server.server_type)
        file.state = KnowledgeFileStateEnum.failed
        await file.save()
        return
    # 动态实例化模型方法
    print('动态实例化方法', embedding_utils.CONFIG.get(embedding_server.server_type))
    config = {
        'api_key': embedding_server.api_key
    }
    embedding_object = getattr(embedding_utils, embedding_utils.CONFIG.get(embedding_server.server_type))(**config)
    # 为了方便进行数据解析处理，此处将去切分后的结果进行分组处理
    all_embedding_info = []
    for i in range(0, len(content), 10):
        # embedding_object.add_text(texts=content[i:i + 10])
        embedding_data = await embedding_object.batch_embedding(
            base_url=embedding_server.server_base_url, text_list=content[i:i + 10]
        )
        if embedding_data:
            all_embedding_info += embedding_data
        else:
            file.state = KnowledgeFileStateEnum.failed
            await file.save()
            return
    # 准备数据存储
    # TODO 当前版本不做动态向量库数据存储
    # 查看配置表中是否存在
    if not embedding_server.table_name:
        # 说明首次使用，需要创建表
        embedding_server.table_name = f'{int(time.time())}'
        await embedding_server.save()
    table_name = embedding_server.table_name
    # 创建Milvus服务对象
    milvus_object = MilvusUtil(
        address=settings.MILVUS_ADDRESS, port=settings.MILVUS_PORT, user=settings.MILVUS_USER,
        password=settings.MILVUS_PASSWORD, db_name=settings.MILVUS_DB_NAME)
    try:
        milvus_object.create_client()
        # 创建表
        milvus_object.create_collection(collection_name=table_name, dim=embedding_server.dimension)
        # load表
        milvus_object.load_collection(collection_name=table_name)
        # 插入数据
        # 构造数据
        all_data = [
            {
                "vector": value, "title": file.file_name, 'document': content[index], "knowledge_id": knowledge.id,
                'file_id': file.id
            } for index, value in enumerate(all_embedding_info)]
        milvus_object.insert(collection_name=table_name, data=all_data)
        file.state = KnowledgeFileStateEnum.success
    except:
        file.state = KnowledgeFileStateEnum.failed
    finally:
        milvus_object.close()
        await file.save()


def handle_file_(file_id: int, knowledge_id: int, user: User):
    loop = asyncio.get_event_loop()
    loop.create_task(handle_file(file_id=file_id, knowledge_id=knowledge_id, user=user))


async def delete_knowledge_file(knowledge_id: int, table_name: str, file_id: int = None):
    # TODO 当前版本不做动态向量库数据存储
    # 创建Milvus服务对象
    milvus_object = MilvusUtil(
        address=settings.MILVUS_ADDRESS, port=settings.MILVUS_PORT, user=settings.MILVUS_USER,
        password=settings.MILVUS_PASSWORD, db_name=settings.MILVUS_DB_NAME)
    try:
        milvus_object.create_client()
        # load表
        milvus_object.load_collection(collection_name=table_name)
        # 执行数据删除
        filter = 'knowledge_id=={}'.format(knowledge_id)
        if file_id:
            filter += ' and file_id=={}'.format(file_id)
        milvus_object.filter_delete(collection_name=table_name, filter=filter)
    except Exception as e:
        print('删除向量失败，', e)
    finally:
        milvus_object.close()
    pass
